import os
import argparse
from docx import Document
from docx.shared import Pt
from docx.enum.text import WD_ALIGN_PARAGRAPH
import pandas as pd

def extract_text_from_word(input_dir, output_file, recursive=False):
    """提取多个Word文档的文本内容"""
    try:
        # 收集所有Word文件
        doc_files = []
        for root, _, files in os.walk(input_dir):
            for file in files:
                if file.lower().endswith('.docx'):
                    doc_files.append(os.path.join(root, file))
            if not recursive:
                break
        
        if not doc_files:
            print("未找到Word文档")
            return
        
        # 提取文本
        with open(output_file, 'w', encoding='utf-8') as out_f:
            for doc_path in doc_files:
                try:
                    doc = Document(doc_path)
                    out_f.write(f"=== {os.path.basename(doc_path)} ===\n\n")
                    
                    for para in doc.paragraphs:
                        if para.text.strip():
                            out_f.write(para.text + "\n")
                    
                    out_f.write("\n" + "="*50 + "\n\n")
                    print(f"已提取: {os.path.basename(doc_path)}")
                except Exception as e:
                    print(f"处理 {doc_path} 失败: {str(e)}")
        
        print(f"\n文本提取完成，保存到: {output_file}")
        
    except Exception as e:
        print(f"提取失败: {str(e)}")

def replace_text_in_word(input_dir, output_dir, search_str, replace_str, recursive=False):
    """批量替换Word文档中的文本"""
    try:
        # 创建输出目录
        os.makedirs(output_dir, exist_ok=True)
        
        # 处理所有Word文件
        doc_files = []
        for root, _, files in os.walk(input_dir):
            for file in files:
                if file.lower().endswith('.docx'):
                    doc_files.append(os.path.join(root, file))
            if not recursive:
                break
        
        if not doc_files:
            print("未找到Word文档")
            return
        
        # 替换文本
        for doc_path in doc_files:
            try:
                doc = Document(doc_path)
                
                # 替换段落中的文本
                for para in doc.paragraphs:
                    if search_str in para.text:
                        para.text = para.text.replace(search_str, replace_str)
                
                # 替换表格中的文本
                for table in doc.tables:
                    for row in table.rows:
                        for cell in row.cells:
                            if search_str in cell.text:
                                cell.text = cell.text.replace(search_str, replace_str)
                
                # 保存新文档
                output_path = os.path.join(output_dir, os.path.basename(doc_path))
                doc.save(output_path)
                print(f"已替换: {os.path.basename(doc_path)}")
            except Exception as e:
                print(f"处理 {doc_path} 失败: {str(e)}")
        
        print(f"\n文本替换完成，结果保存到: {output_dir}")
        
    except Exception as e:
        print(f"替换失败: {str(e)}")

def generate_from_template(template_path, data_file, output_dir):
    """根据模板和数据批量生成Word文档"""
    try:
        # 读取数据
        ext = os.path.splitext(data_file)[1].lower()
        if ext == '.csv':
            df = pd.read_csv(data_file)
        else:
            df = pd.read_excel(data_file)
        
        # 创建输出目录
        os.makedirs(output_dir, exist_ok=True)
        
        # 逐个生成文档
        for _, row in df.iterrows():
            try:
                doc = Document(template_path)
                # 替换占位符（格式为{{column_name}}）
                for para in doc.paragraphs:
                    for key, value in row.items():
                        placeholder = f"{{{{{key}}}}}"
                        if placeholder in para.text:
                            para.text = para.text.replace(placeholder, str(value))
                
                # 处理表格中的占位符
                for table in doc.tables:
                    for row_table in table.rows:
                        for cell in row_table.cells:
                            for key, value in row.items():
                                placeholder = f"{{{{{key}}}}}"
                                if placeholder in cell.text:
                                    cell.text = cell.text.replace(placeholder, str(value))
                
                # 生成文件名（使用第一列作为标识）
                filename = f"{row.iloc[0]}.docx"
                output_path = os.path.join(output_dir, filename)
                doc.save(output_path)
                print(f"已生成: {filename}")
            except Exception as e:
                print(f"生成文档失败: {str(e)}")
        
        print(f"\n文档生成完成，共 {len(df)} 个文件，保存到: {output_dir}")
        
    except Exception as e:
        print(f"生成失败: {str(e)}")

if __name__ == "__main__":
    # 检查依赖
    try:
        from docx import Document
    except ImportError:
        print("请安装依赖: pip install python-docx")
        exit(1)
    
    parser = argparse.ArgumentParser(description='Word文档处理工具')
    parser.add_argument('operation', choices=['extract', 'replace', 'generate'], 
                      help='操作类型：extract(提取文本)、replace(替换文本)、generate(从模板生成)')
    parser.add_argument('input', help='输入目录或文件')
    parser.add_argument('output', help='输出目录或文件')
    parser.add_argument('--search', help='替换时的搜索文本')
    parser.add_argument('--replace', help='替换时的替换文本')
    parser.add_argument('--template', help='生成文档时的模板文件')
    parser.add_argument('--recursive', action='store_true', help='递归处理子目录')
    
    args = parser.parse_args()
    
    if args.operation == 'extract':
        extract_text_from_word(args.input, args.output, args.recursive)
    elif args.operation == 'replace':
        if not args.search or not args.replace:
            print("请指定--search和--replace参数")
        else:
            replace_text_in_word(args.input, args.output, args.search, args.replace, args.recursive)
    elif args.operation == 'generate':
        if not args.template:
            print("请指定--template参数")
        else:
            generate_from_template(args.template, args.input, args.output)
